feat: mmap lazy loading and phase-gated VRAM swap - #44
Conversation
Replaced fread pipeline with a cross-platform memory-mapped file/mmap. This reduces boot times by allowing lazy weight loading for sub-models and no copy DMA. - Introducing MappedFile RAII wrapper for zero-copy memory-mapped file I/O supporting POSIX (mmap/madvise) and Windows (CreateFileMapping). - Decoupled GGUF metadata parsing from VRAM/RAM buffer allocation. - Audio Codec consumes 0 MB VRAM at init, Weight buffers are allocated on-demand, VQ codebook caches are populated directly from the mmap pointer into system RAM. - Slow-AR model weights are page-faulted from the mmap pointer to VRAM, bypassing intermediate buffers. - We've replaced read_all_tensor_data and read_tensor_data in favor of lazy loading.
… for server mode Utilizes mmap and lazy loading introduced in the previous commit to dynamically manage VRAM occupancy. Intelligently swapping in and out the required submodels depending on which phase of the processing we're at. This minimizes both peak and idle VRAM usage, allowing running larger models without OOM and increases speeds by reducing memory pressure. - Phase-Gated VRAM Swapping: Slow-AR weights and KV cache are freed immediately after generation completes, right before Audio Codec weights are restored for decode. .. We don't need the 4.2 GB (Q8_0) SlowAR to occupy VRAM as we're running inference on the Audio Codec part and possibly crash via OOM. .. Without this system, Q8_0 would hit 7-7.5 GB VRAM usage during final phase of the processing (Audio Codec) in one sentence long generation. Now it's just ~2.3 GB (Vulkan, Linux latest MESA) - CLI flags --no-vram-swap (opt out) and --hot-swap (opt in) to customize behavior. - vram-swap retains the OS page cache between requests, pagefaulting we read from RAM instead of disk, this only takes a few seconds. .. Also keeps compute buffers etc. in VRAM which are relatively small (~168 MB Vulkan) so we can immediately begin processing. .. The gguf occupies system RAM instead of VRAM, however, this occupancy is not "locked" meaning OS will free it for other applications as needed. .. This is basically tells the OS "Here's this memory pool that maps to compute buffer, keep it alive but also don't hesitate to free the memory if other apps need it." - Aggressive hot-swap mode goes a step further and explicitly instructs the kernel to reclaim memory. .. This is optimal if you want minimal, 100 MB RAM + 25 MB VRAM idles without bothering the OS and have the model on flash storage. - Backend synchronization via ggml_backend_synchronize to ensure different backends (Vulkan, CUDA, Metal, etc.) reclaim memory sooner than later to prevent PCIe thrashing. .. This is critical to reduce peak VRAM usage therefore allowing us to run larger quants without filling VRAM to the brim. Also reduces pressure on other apps and their VRAM occupancies. - Background prefetching - spawns a background thread to restore Slow-AR weights concurrently with CPU-bound voice profile loading to hide PCIe latency. - Thread safe, all synthesis entry points join pending_offload_thread_ before proceeding to prevent race conditions between background eviction and new weight restoration.
📝 WalkthroughWalkthroughThe PR adds cross-platform mmap support for GGUF files, defers model and codec weight allocation, exposes GPU residency controls, and coordinates VRAM swapping through pipeline configuration, CLI flags, and background offload threads. ChangesMapped weight lifecycle and VRAM hot-swapping
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Pipeline
participant SlowARModel
participant AudioCodec
participant MappedFile
Pipeline->>SlowARModel: restore_weights_to_gpu()
SlowARModel->>MappedFile: read mapped tensor data
Pipeline->>AudioCodec: restore_weights_to_gpu()
AudioCodec->>MappedFile: read mapped tensor data
Pipeline->>SlowARModel: free_gpu_weights()
Pipeline->>AudioCodec: free_gpu_weights()
Pipeline->>MappedFile: drop_page_cache()
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (6)
src/s2_codec.cpp (2)
942-942: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBinding a
const&to a conditional with a temporary — unnecessary full-set copy (cppcheckdanglingTemporaryLifetime). Because the?:yields a prvalue, theModel->weight_tensor_set()branch is materialized into a temporary and the whole set is copied even whenModel != nullptr; the reference then binds to that (lifetime-extended) temporary. It is not actually dangling, but the copy is wasteful. Prefer a pointer to avoid the copy and silence the analyzer.♻️ Use a pointer instead of a copied reference
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + static const std::unordered_set<ggml_tensor*> empty_weights; + const auto & model_weights = Model ? Model->weight_tensor_set() : empty_weights;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_codec.cpp` at line 942, Update the model_weights initialization in the surrounding codec flow to use a pointer to Model->weight_tensor_set() when Model is non-null, avoiding the conditional’s temporary full-set copy; provide an appropriate empty-set fallback for the null-Model case and adjust subsequent accesses to dereference the pointer while preserving existing behavior.Source: Linters/SAST tools
129-155: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy lift
allocate_codec_buffersshould respect per-allocation caps. On backends like Vulkan, one largeggml_backend_buft_alloc_buffer()can fail once the request exceeds the backend’s max allocation size, even when total VRAM is available. Chunk this likeallocate_weight_buffersdoes, or guard againstggml_backend_buft_get_max_size()before a single allocation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_codec.cpp` around lines 129 - 155, Update allocate_codec_buffers to respect the buffer type’s maximum allocation size, using ggml_backend_buft_get_max_size() and the chunking strategy established by allocate_weight_buffers. Ensure no single ggml_backend_buft_alloc_buffer call exceeds that cap while preserving alignment, total-byte accounting, error reporting, and output-buffer initialization.src/s2_model.cpp (1)
1194-1226: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRestore path re-allocates and re-copies CPU weights unnecessarily.
free_gpu_weightsfrees onlymodel_bufs_gpu, butrestore_weights_to_gpuresetsweights_allocated_and callsallocate_and_load_weights, which unconditionally re-runs the CPU branch (allocate_weight_buffers(backend_cpu_, ...)frees and re-places all CPU tensors and re-copies them from mmap). Consider guarding the CPU allocation so only GPU weights are re-materialized on restore.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_model.cpp` around lines 1194 - 1226, Update allocate_and_load_weights so the restore path does not reallocate or recopy CPU weights when they remain valid; guard the backend_cpu_ allocation and CPU tensor loading using the same GPU-restore state or condition that indicates CPU weights are already materialized, while preserving initial CPU allocation and loading behavior.src/s2_pipeline.cpp (2)
732-754: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBackground pre-fetch thread duplicated with the streaming path.
This block (condition, thread body, join-on-failure/success) is duplicated near-verbatim in
synthesize_streaming_raw(Lines 978-1002). Extracting a small helper (e.g.start_background_model_prefetch()/join_background_model_prefetch()) would avoid future logic drift between the two call sites.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_pipeline.cpp` around lines 732 - 754, Extract the duplicated VRAM pre-fetch lifecycle from the current function and synthesize_streaming_raw into shared helpers such as start_background_model_prefetch and join_background_model_prefetch. Preserve the existing enable_vram_swap/is_persistent/model_prefers_gpu_/is_weights_on_gpu condition, background acquire_compute_resources/restore_weights_to_gpu work, and joins on both failure and success paths.
893-916: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHot-swap cleanup thread duplicated with the streaming path.
This entire persistent/hot-swap cleanup block (compute-buffer release + background offload thread that frees GPU weights and drops page caches) is duplicated near-verbatim in
synthesize_streaming_prompt_codes_locked(Lines 1264-1288). Same drift risk as the pre-fetch duplication above; consider a shared private helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_pipeline.cpp` around lines 893 - 916, Extract the duplicated persistent hot-swap cleanup logic from the current block and synthesize_streaming_prompt_codes_locked into a shared private helper. The helper must release compute buffers, asynchronously free model and codec GPU weights, drop both mapped-file page caches, and update pending_offload_thread_; call it from both paths while preserving existing non-hot-swap and single-shot behavior.src/main.cpp (1)
93-94: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
--hot-swapcan silently no-op depending on flag combinations.
--hot-swapcleanup is gated onparams.is_persistent(only settruefor--server, Line 316), and the entire hot-swap block ins2_pipeline.cppis additionally nested underparams.enable_vram_swap. So--hot-swapalone in CLI/file mode does nothing, and--hot-swap --no-vram-swaptogether also silently disables hot-swap. Consider warning the user in these cases, similar to the existing warnings pattern (e.g. Lines 271-277).💡 Suggested warning
+ if (params.enable_hot_swap && !use_server) { + safe_print_error_ln("Warning: --hot-swap has no effect outside --server mode.\n"); + } + if (params.enable_hot_swap && !params.enable_vram_swap) { + safe_print_error_ln("Warning: --hot-swap has no effect when combined with --no-vram-swap.\n"); + }Also applies to: 191-202
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/main.cpp` around lines 93 - 94, Add validation warnings in the argument/configuration handling near the existing warnings and the --hot-swap option: warn when --hot-swap is enabled without persistent server mode, and when it is combined with --no-vram-swap, since the s2_pipeline.cpp cleanup path requires both conditions. Keep the existing behavior unchanged while clearly informing users that hot-swap will be inactive.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/s2_codec.cpp`:
- Around line 934-936: Update reset_codec_impl to release impl_->backend_cpu
during reset and destruction, matching the existing cleanup for impl_->backend,
impl_->ctx_w, and impl_->model_buf. Ensure the pointer is cleared after freeing,
and preserve the current initialization behavior in the backend_cpu setup block.
- Around line 1528-1544: Update read_f32 in refresh_host_caches_from_mmap to
throw an error for tensor types other than GGML_TYPE_F32 and GGML_TYPE_F16,
preventing unsupported data from being returned as zero-filled values. Also
update the base calculation in the same function to widen c/code before
multiplication, avoiding int32_t overflow while preserving the existing indexing
behavior.
In `@src/s2_pipeline.cpp`:
- Around line 442-455: Update the codec state assignment in the initialization
flow so codec_prefers_gpu_ reflects the codec’s actual post-load backend, not
the pre-fallback use_gpu_codec request. Reuse the codec’s existing GPU-residency
query or equivalent state established after the GPU load/fallback logic, then
keep the VRAM State Machine diagnostics and downstream swap decisions based on
that corrected flag.
---
Nitpick comments:
In `@src/main.cpp`:
- Around line 93-94: Add validation warnings in the argument/configuration
handling near the existing warnings and the --hot-swap option: warn when
--hot-swap is enabled without persistent server mode, and when it is combined
with --no-vram-swap, since the s2_pipeline.cpp cleanup path requires both
conditions. Keep the existing behavior unchanged while clearly informing users
that hot-swap will be inactive.
In `@src/s2_codec.cpp`:
- Line 942: Update the model_weights initialization in the surrounding codec
flow to use a pointer to Model->weight_tensor_set() when Model is non-null,
avoiding the conditional’s temporary full-set copy; provide an appropriate
empty-set fallback for the null-Model case and adjust subsequent accesses to
dereference the pointer while preserving existing behavior.
- Around line 129-155: Update allocate_codec_buffers to respect the buffer
type’s maximum allocation size, using ggml_backend_buft_get_max_size() and the
chunking strategy established by allocate_weight_buffers. Ensure no single
ggml_backend_buft_alloc_buffer call exceeds that cap while preserving alignment,
total-byte accounting, error reporting, and output-buffer initialization.
In `@src/s2_model.cpp`:
- Around line 1194-1226: Update allocate_and_load_weights so the restore path
does not reallocate or recopy CPU weights when they remain valid; guard the
backend_cpu_ allocation and CPU tensor loading using the same GPU-restore state
or condition that indicates CPU weights are already materialized, while
preserving initial CPU allocation and loading behavior.
In `@src/s2_pipeline.cpp`:
- Around line 732-754: Extract the duplicated VRAM pre-fetch lifecycle from the
current function and synthesize_streaming_raw into shared helpers such as
start_background_model_prefetch and join_background_model_prefetch. Preserve the
existing enable_vram_swap/is_persistent/model_prefers_gpu_/is_weights_on_gpu
condition, background acquire_compute_resources/restore_weights_to_gpu work, and
joins on both failure and success paths.
- Around line 893-916: Extract the duplicated persistent hot-swap cleanup logic
from the current block and synthesize_streaming_prompt_codes_locked into a
shared private helper. The helper must release compute buffers, asynchronously
free model and codec GPU weights, drop both mapped-file page caches, and update
pending_offload_thread_; call it from both paths while preserving existing
non-hot-swap and single-shot behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bc17f0ca-6456-4aec-9e41-64291ff19de0
📒 Files selected for processing (10)
CMakeLists.txtinclude/s2_codec.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cpp
| if (!impl_->backend_cpu) { | ||
| impl_->backend_cpu = ggml_backend_cpu_init(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
backend_cpu is initialized but never freed (and appears unused). reset_codec_impl (Lines 193-208) frees impl.backend, impl.ctx_w, and impl.model_buf, but not impl.backend_cpu, so this ggml_backend_cpu_init() leaks on every reset/destruction. It also does not appear to be used by the allocation/compute paths (those use impl_->backend). Either wire backend_cpu into the lifecycle and free it in reset_codec_impl, or drop the field.
🛠️ Free backend_cpu in reset_codec_impl
if (impl.backend) {
ggml_backend_free(impl.backend);
impl.backend = nullptr;
}
+ if (impl.backend_cpu) {
+ ggml_backend_free(impl.backend_cpu);
+ impl.backend_cpu = nullptr;
+ }
impl = AudioCodec::Impl();🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_codec.cpp` around lines 934 - 936, Update reset_codec_impl to release
impl_->backend_cpu during reset and destruction, matching the existing cleanup
for impl_->backend, impl_->ctx_w, and impl_->model_buf. Ensure the pointer is
cleared after freeing, and preserve the current initialization behavior in the
backend_cpu setup block.
| bool AudioCodec::refresh_host_caches_from_mmap() { | ||
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | ||
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | ||
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | ||
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | ||
| auto it = impl_->tensor_offsets.find(t); | ||
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | ||
| const size_t n = ggml_nelements(t); | ||
| std::vector<float> out(n); | ||
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | ||
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | ||
| else if (t->type == GGML_TYPE_F16) { | ||
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | ||
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | ||
| } | ||
| return out; | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
refresh_host_caches_from_mmap silently zero-fills unsupported tensor types. read_f32 handles only GGML_TYPE_F32/GGML_TYPE_F16; any other type falls through with out left zero-initialized, producing silently-wrong VQ codebooks instead of a hard failure. The prior tensor_to_f32 threw on unsupported types. Also, base = c * cb_dim (Line 1553) multiplies two int32_t before widening to size_t, risking overflow for large codebooks; cast first as done elsewhere (static_cast<size_t>(code) * codebook_dim).
🛡️ Fail loudly on unsupported types
if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float));
else if (t->type == GGML_TYPE_F16) {
const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src);
for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]);
- }
+ } else {
+ throw std::runtime_error("unsupported vq tensor type: " + name);
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| bool AudioCodec::refresh_host_caches_from_mmap() { | |
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | |
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | |
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | |
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | |
| auto it = impl_->tensor_offsets.find(t); | |
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | |
| const size_t n = ggml_nelements(t); | |
| std::vector<float> out(n); | |
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | |
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | |
| else if (t->type == GGML_TYPE_F16) { | |
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | |
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | |
| } | |
| return out; | |
| }; | |
| bool AudioCodec::refresh_host_caches_from_mmap() { | |
| if (!impl_ || !impl_->mapped_gguf_.is_open()) return false; | |
| auto read_f32 = [&](const std::string& name) -> std::vector<float> { | |
| ggml_tensor* t = ggml_get_tensor(impl_->ctx_w, name.c_str()); | |
| if (!t) throw std::runtime_error("missing vq tensor: " + name); | |
| auto it = impl_->tensor_offsets.find(t); | |
| if (it == impl_->tensor_offsets.end()) throw std::runtime_error("missing offset"); | |
| const size_t n = ggml_nelements(t); | |
| std::vector<float> out(n); | |
| const uint8_t* src = impl_->mapped_gguf_.data() + impl_->gguf_data_offset + it->second; | |
| if (t->type == GGML_TYPE_F32) std::memcpy(out.data(), src, n * sizeof(float)); | |
| else if (t->type == GGML_TYPE_F16) { | |
| const ggml_fp16_t* tmp = reinterpret_cast<const ggml_fp16_t*>(src); | |
| for (size_t i = 0; i < n; ++i) out[i] = ggml_fp16_to_fp32(tmp[i]); | |
| } else { | |
| throw std::runtime_error("unsupported vq tensor type: " + name); | |
| } | |
| return out; | |
| }; |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_codec.cpp` around lines 1528 - 1544, Update read_f32 in
refresh_host_caches_from_mmap to throw an error for tensor types other than
GGML_TYPE_F32 and GGML_TYPE_F16, preventing unsupported data from being returned
as zero-filled values. Also update the base calculation in the same function to
widen c/code before multiplication, avoiding int32_t overflow while preserving
the existing indexing behavior.
| initialized_ = true; | ||
|
|
||
| model_prefers_gpu_ = model().is_weights_on_gpu(); | ||
| codec_prefers_gpu_ = use_gpu_codec; | ||
|
|
||
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | ||
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | ||
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | ||
| } else { | ||
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
codec_prefers_gpu_ records intent, not the actual post-fallback backend.
model_prefers_gpu_ is derived from the actual result via model().is_weights_on_gpu() (Line 444), but codec_prefers_gpu_ is assigned use_gpu_codec (Line 445) — a value computed before the GPU load attempt/fallback at Lines 391-417. If the codec's GPU load fails and it falls back to CPU (Lines 402-409), codec_prefers_gpu_ stays true even though the codec is actually CPU-backed. This corrupts the "VRAM State Machine" diagnostics (wrong case printed) and feeds every downstream swap decision (Lines 815, 821, 859, 873, 911, 1045, 1283) with a flag that doesn't reflect reality. It happens to be masked today only because AudioCodec::restore_weights_to_gpu()/free_gpu_weights() no-op when there are no GPU weights, but the state is still wrong and fragile.
🐛 Proposed fix
- codec_prefers_gpu_ = use_gpu_codec;
+ codec_prefers_gpu_ = codec().is_weights_on_gpu();📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| initialized_ = true; | |
| model_prefers_gpu_ = model().is_weights_on_gpu(); | |
| codec_prefers_gpu_ = use_gpu_codec; | |
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | |
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | |
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | |
| } else { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | |
| } | |
| initialized_ = true; | |
| model_prefers_gpu_ = model().is_weights_on_gpu(); | |
| codec_prefers_gpu_ = codec().is_weights_on_gpu(); | |
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 1 (Both prefer GPU) - Codec is lazily allocated on demand."); | |
| } else if (model_prefers_gpu_ && !codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 2 (Slow-AR GPU, Codec CPU) - Ready."); | |
| } else if (!model_prefers_gpu_ && codec_prefers_gpu_) { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 3 (Slow-AR CPU, Codec GPU) - Codec is lazily allocated on demand."); | |
| } else { | |
| safe_print_ln("[Pipeline] VRAM State Machine: Case 4 (All CPU) - Ready."); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 442 - 455, Update the codec state
assignment in the initialization flow so codec_prefers_gpu_ reflects the codec’s
actual post-load backend, not the pre-fallback use_gpu_codec request. Reuse the
codec’s existing GPU-residency query or equivalent state established after the
GPU load/fallback logic, then keep the VRAM State Machine diagnostics and
downstream swap decisions based on that corrected flag.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/s2_codec.cpp (1)
942-942: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAvoid copying the entire model weight set.
The ternary mixes an lvalue reference (
Model->weight_tensor_set()) with a prvalue temporary, so the result is a prvalue andmodel_weightsbinds to a full copy of the set (this is also why cppcheck flags line 955 — it's a false positive since lifetime extension applies). Prefer a pointer to avoid the copy; it also makes theModel &&guard meaningful.♻️ Suggested change
- const auto & model_weights = Model ? Model->weight_tensor_set() : std::unordered_set<ggml_tensor*>(); + const std::unordered_set<ggml_tensor*> * model_weights = + Model ? &Model->weight_tensor_set() : nullptr;and at line 955:
- if (Model && model_weights.find(t) != model_weights.end()) continue; + if (model_weights && model_weights->find(t) != model_weights->end()) continue;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_codec.cpp` at line 942, Update the model_weights initialization near weight_tensor_set() to use a pointer, selecting the model’s weight set only when Model is non-null and otherwise storing nullptr. Adjust its use near line 955 to dereference the pointer only under the existing Model guard, preserving lifetime safety without copying the entire set.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/s2_mapped_file.cpp`:
- Around line 158-171: Update MappedFile::drop_page_cache so the madvise call
remains available on both Linux and macOS, while the fd_ posix_fadvise call and
POSIX_FADV_DONTNEED usage are compiled only under __linux__. Keep the existing
Windows VirtualUnlock branch unchanged.
In `@src/s2_pipeline.cpp`:
- Around line 895-909: Reorder the offline hot-swap flow around the
pending_offload_thread_ creation: perform clear_kv_cache(), Post-Phase3
diagnostics, and all model()/codec() GPU-memory or residency reads before
starting the background thread. Spawn the offload_thread only as the final
action so its free_gpu_weights() and drop_page_cache() calls cannot race with
remaining state access.
---
Nitpick comments:
In `@src/s2_codec.cpp`:
- Line 942: Update the model_weights initialization near weight_tensor_set() to
use a pointer, selecting the model’s weight set only when Model is non-null and
otherwise storing nullptr. Adjust its use near line 955 to dereference the
pointer only under the existing Model guard, preserving lifetime safety without
copying the entire set.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bc17f0ca-6456-4aec-9e41-64291ff19de0
📒 Files selected for processing (10)
CMakeLists.txtinclude/s2_codec.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cpp
| if (params.enable_hot_swap) { | ||
| safe_print_ln("[Pipeline] Hot-Swap: Releasing compute resources..."); | ||
| model().free_compute_buffers(); | ||
|
|
||
| safe_print_ln("[Pipeline] Hot-Swap: Spawning background thread to free VRAM & RAM..."); | ||
| std::thread offload_thread([this]() { | ||
| if (model().is_weights_on_gpu()) model().free_gpu_weights(); | ||
| if (codec().is_weights_on_gpu()) codec().free_gpu_weights(); | ||
|
|
||
| model().mapped_file().drop_page_cache(); | ||
| codec().mapped_file().drop_page_cache(); | ||
|
|
||
| safe_print_ln("[Pipeline] Hot-Swap: Background VRAM & RAM free complete."); | ||
| }); | ||
| pending_offload_thread_ = std::move(offload_thread); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Data race: the hot-swap offload thread frees GPU weights while the main thread still reads them.
In the offline persistent hot-swap path the background thread is spawned here and immediately begins model().free_gpu_weights() / codec().free_gpu_weights() (both call free_backend_buffers(...), clearing weights_.model_bufs_gpu and resetting residency flags). But the main thread keeps running after the spawn and reads the very same shared state:
- Line 923:
model().clear_kv_cache() - Lines 955–959:
model().get_gpu_memory_usage_bytes()/codec().get_gpu_memory_usage_bytes(), which iteratemodel_bufs_gpu/ readmodel_bufconcurrently with the thread clearing them.
This is undefined behavior (iterating a vector being cleared, reading freed buffer sizes) and can crash. The streaming variant is safe because it spawns the offload thread after all diagnostics; the offline path emits Post-Phase3 after the spawn. Move the post-decode diagnostics and any model/codec state reads before spawning the offload thread (or snapshot the values first).
🔒 Suggested direction
Emit the Post-Phase3 VRAM diagnostic and finish all get_gpu_memory_usage_bytes() reads (and clear_kv_cache()) before line 900, then spawn offload_thread as the last action so the background thread has exclusive access to the weight/buffer state it mutates.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 895 - 909, Reorder the offline hot-swap
flow around the pending_offload_thread_ creation: perform clear_kv_cache(),
Post-Phase3 diagnostics, and all model()/codec() GPU-memory or residency reads
before starting the background thread. Spawn the offload_thread only as the
final action so its free_gpu_weights() and drop_page_cache() calls cannot race
with remaining state access.
Replace MADV_RANDOM with MADV_SEQUENTIAL. Weight loading iterates tensors in sequential file order, so disabling readahead forced ~1.3M individual 4 KB I/O syscalls on ..cold reads after drop_page_cache(). MADV_SEQUENTIAL enables aggressive kernel readahead from byte 0, reducing syscall count. - Add MADV_SEQUENTIAL (Linux+macOS), FILE_FLAG_SEQUENTIAL_SCAN (Windows) as the equivalent hint. - Add MADV_HUGEPAGE (Linux) to reduce TLB pressure during multi-GB loads - Add MADV_DONTDUMP (Linux) to exclude the mapping from core dumps
…ipeline Extended the phase-gated VRAM swap with finer-grained codec weight management, and concurrent decode threading. - Codec encoder/decoder granular split: codec weights are classified at load time into encoder and decoder groups via tensor name prefixes. - New methods (free/restore/is_on_gpu/get_bytes for each group) allowing the pipeline to load only the encoder for reference audio encoding, free it before generation, ..load only the decoder for the decode phase. - Minimizing VRAM footprint at each step while catering to lower latency by utilizing background threads to lazy load as needed. - The priority is to reduce Slow-AR processing stage's VRAM footprint as this is the heaviest model in the pipeline, so we never keep another (unnecessary) model loaded when that's on. And we free it before loading any new models. In effect, OOM is far less likely and as long as your GPU can fit the Slow-AR and its buffers it'll run the whole pipeline without issue. - Deferred weight loading: when VRAM swap is active and the model prefers GPU, init() skips Slow-AR weight allocation entirely and calls warm_page_cache() to pre-fault mmap pages via ..MADV_WILLNEED + MADV_COLD (Linux), fcntl F_RDAHEAD (macOS), or PrefetchVirtualMemory ..(Windows). First-request restore hits warm RAM instead of cold disk. Replacing the old incorrect VirtualUnlock. - prefers_gpu() provides an intent-based check that works before weights are loaded, replacing is_weights_on_gpu() for init decisions to handle edge cases better. - Codec eviction + Slow-AR restore and KV cache init now runs in background threads, hiding PCIe latency behind the CPU bound prompt construction. - Overlapped decode path: when Slow-AR is on GPU and codec is on CPU, a producer-consumer thread pair decodes audio frames concurrently with generation via mutex, condition variable, and atomic frame counters. This reduces Total RTF by starting the CPU codec work as early as we can instead of waiting for GPU to finish Slow-AR processing in its entirety. - server-aware swapping: more_segments_pending server sets this flag on all sentence segments except the last within a single request, keeping Slow-AR resident in VRAM across segments and eliminating per-segment restore overhead. - Streaming path: granular codec management (free encoder, restore decoder only), Slow-AR freed in non-hot-swap persistent mode (fixes VRAM leak where Slow-AR was never freed after streaming requests). - pre_restore_thread removed from synthesize_raw/synthesize_streaming_raw; replaced with pending_offload_thread_ join before encoding to prevent races between background eviction and encoder weight restoration. - --fast-decoder-cpu / --codebook-cpu CLI flags force specific tensor groups onto CPU, saving ~200-400 MB / ~56 MB VRAM respectively at the cost of PCIe transfers. Previously, any --gpu-layers value also offloaded fast-decoder and codebook tensors. These flags decouple that decision for finer-grained VRAM control. - allocate_weight_buffers nulls stale tensor data/buffer pointers after freeing, preventing use-after-free when weights are re-allocated after a free/restore cycle. - MappedFile::open uses CreateFileW with UTF-8 -> UTF-16 conversion, fixing model loading on Windows paths containing non-ASCII characters.
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/s2_model.cpp (2)
1200-1212: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
get_gpu_memory_usage_bytescounts the KV cache even when it lives on the CPU.
init_kv_cacheline 624 allocateskv_buf_onbackend_cpu_whenn_gpu_layers_ == 0orbackend_gpu_is null. This function addskv_buf_unconditionally, so a CPU-only model reports nonzero GPU usage.The pipeline prints this value in every
[VRAM Diag]line, so the CPU-only and hybrid configurations report inflated VRAM figures.🐛 Proposed fix
- if (kv_buf_) { + const bool kv_on_gpu = (n_gpu_layers_ > 0 && backend_gpu_ != nullptr); + if (kv_buf_ && kv_on_gpu) { total += ggml_backend_buffer_get_size(kv_buf_); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_model.cpp` around lines 1200 - 1212, Update SlowARModel::get_gpu_memory_usage_bytes so kv_buf_ contributes to the total only when the KV cache is allocated on the GPU; preserve the existing weight-buffer accounting and exclude CPU-backed KV caches used by CPU-only configurations.
1230-1240: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winBound the mmap reads and do not silently skip tensors without an offset.
Two problems in the copy loops:
No bounds check.
base + gguf_data_offset_ + it->secondplusggml_nbytes(t)is read with no comparison againstmapped_gguf_.size(). The offsets come from GGUF metadata, but the mapping is a second, independentopen()of the same path (line 559). If the file is truncated, replaced, or malformed, the read runs past the end of the mapping. On Linux a read past the last mapped page of a shortened file raises SIGBUS, which noboolreturn can recover from.A weight tensor that has no entry in
tensor_offsets_is skipped without an error. Its buffer keeps whatever the allocator left in it, so the model runs on uninitialized weights and produces garbage output with no diagnostic.🛡️ Proposed fix
const uint8_t* base = mapped_gguf_.data(); - for (ggml_tensor * t : original_gpu_weights_) { - auto it = tensor_offsets_.find(t); - if (it != tensor_offsets_.end()) - ggml_backend_tensor_set(t, base + gguf_data_offset_ + it->second, 0, ggml_nbytes(t)); - } - for (ggml_tensor * t : original_cpu_weights_) { - auto it = tensor_offsets_.find(t); - if (it != tensor_offsets_.end()) - ggml_backend_tensor_set(t, base + gguf_data_offset_ + it->second, 0, ggml_nbytes(t)); - } + const size_t mapped_size = mapped_gguf_.size(); + + auto load_tensor = [&](ggml_tensor * t) -> bool { + auto it = tensor_offsets_.find(t); + if (it == tensor_offsets_.end()) { + std::cerr << "[Model] missing GGUF offset for tensor '" << t->name << "'" << std::endl; + return false; + } + const size_t nbytes = ggml_nbytes(t); + const size_t begin = gguf_data_offset_ + it->second; + if (begin > mapped_size || nbytes > mapped_size - begin) { + std::cerr << "[Model] tensor '" << t->name << "' extends past the mapped file" + << std::endl; + return false; + } + ggml_backend_tensor_set(t, base + begin, 0, nbytes); + return true; + }; + + for (ggml_tensor * t : original_gpu_weights_) { + if (!load_tensor(t)) return false; + } + for (ggml_tensor * t : original_cpu_weights_) { + if (!load_tensor(t)) return false; + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/s2_model.cpp` around lines 1230 - 1240, Update the tensor copy loops around original_gpu_weights_ and original_cpu_weights_ to validate each tensor’s offset and byte range against mapped_gguf_.size() before calling ggml_backend_tensor_set; reject overflow and out-of-bounds ranges through the surrounding load/error path so no mmap read occurs. Also treat a missing tensor_offsets_ entry as an explicit load failure instead of silently skipping the tensor, preserving diagnostics for both GPU and CPU weights.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/s2_mapped_file.cpp`:
- Around line 204-210: Remove the conditional MADV_COLD madvise call from
warm_page_cache(), leaving only the page-warming advice there. If cold-page
reclamation is required, apply MADV_COLD in drop_page_cache() instead.
In `@src/s2_pipeline.cpp`:
- Around line 546-552: Weight-restore results are ignored on three VRAM-swap
paths. In src/s2_pipeline.cpp lines 546-552, check
codec().restore_encoder_weights() and return false before encode; in lines
1083-1089, check codec().restore_decoder_weights() and fail before
decode_codes_windowed; in lines 1243-1256, check both
model().restore_weights_to_gpu() and codec().restore_decoder_weights(), report
failures through sink.on_error(...), and return false, matching the offline
restore handling.
- Around line 436-446: Update synthesize_prompt_codes_locked and
synthesize_streaming_prompt_codes_locked to restore weights based on model
state, not the request’s enable_vram_swap flag: before generation, when
model().is_weights_on_gpu() is false and model_prefers_gpu_ is true, acquire
compute resources and call restore_weights_to_gpu(), returning false with the
existing error-reporting style on failure. Retain the current enable_vram_swap
blocks only for overlap and eviction optimizations.
- Around line 1053-1068: In the hot-swap flow, move the Post-Phase3 diagnostic
and model().clear_kv_cache() before the background thread is created, then make
the offload-thread creation and assignment to pending_offload_thread_ the final
action in the function. Ensure no subsequent code reads model or codec GPU
residency, buffers, or memory usage after the spawn.
- Around line 816-888: Serialize VRAM phase-1 work with KV-cache operations in
the generation flow: move model().clear_kv_cache() before starting
vram_phase1_thread, and ensure vram_phase1_thread joins before spawning
kv_init_thread unless the KV cache is definitively CPU-resident. Remove all
get_gpu_memory_usage_bytes() calls from the phase-1 lambda, then emit the
combined VRAM diagnostic only after both threads have joined, while preserving
vram_phase1_ok and kv_init_ok failure handling.
- Around line 979-1021: Add an RAII thread-joining helper near
CodecDecodeCacheScope, then guard decode_thread, vram_phase1_thread, and
kv_init_thread with it so joinable threads are joined during exceptional exits.
Replace the manual decode_thread.join() with the helper’s cleanup while
preserving existing completion signaling and thread behavior.
---
Outside diff comments:
In `@src/s2_model.cpp`:
- Around line 1200-1212: Update SlowARModel::get_gpu_memory_usage_bytes so
kv_buf_ contributes to the total only when the KV cache is allocated on the GPU;
preserve the existing weight-buffer accounting and exclude CPU-backed KV caches
used by CPU-only configurations.
- Around line 1230-1240: Update the tensor copy loops around
original_gpu_weights_ and original_cpu_weights_ to validate each tensor’s offset
and byte range against mapped_gguf_.size() before calling
ggml_backend_tensor_set; reject overflow and out-of-bounds ranges through the
surrounding load/error path so no mmap read occurs. Also treat a missing
tensor_offsets_ entry as an explicit load failure instead of silently skipping
the tensor, preserving diagnostics for both GPU and CPU weights.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4923cd18-4bd3-46a2-b694-a7f805904e24
📒 Files selected for processing (10)
include/s2_codec.hinclude/s2_mapped_file.hinclude/s2_model.hinclude/s2_pipeline.hsrc/main.cppsrc/s2_codec.cppsrc/s2_mapped_file.cppsrc/s2_model.cppsrc/s2_pipeline.cppsrc/s2_server.cpp
🚧 Files skipped from review as they are similar to previous changes (3)
- src/main.cpp
- include/s2_mapped_file.h
- src/s2_codec.cpp
| #ifdef __linux__ | ||
| ::madvise(data_, size_, MADV_WILLNEED); | ||
| ::madvise(data_, size_, MADV_SEQUENTIAL); | ||
|
|
||
| #ifdef MADV_COLD | ||
| ::madvise(data_, size_, MADV_COLD); | ||
| #endif |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Linux madvise MADV_COLD semantics inactive list reclaim
💡 Result:
The MADV_COLD flag, introduced in Linux 5.4, is a non-destructive advice hint that tells the kernel a specified memory range is not expected to be accessed in the near future [1][2]. Its primary semantic is to "deactivate" pages in that range, making them more probable targets for memory reclamation during periods of system memory pressure [1][3]. Key aspects of its operation include: Deactivation Strategy: MADV_COLD moves pages from the active LRU (Least Recently Used) list to the inactive LRU list [4][3]. By lowering the activation status of these pages, the kernel identifies them as lower-priority for resident memory, allowing the standard memory management subsystem to reclaim them earlier than other active pages if memory pressure arises [4]. Non-destructive Nature: Unlike MADV_DONTNEED (which discards pages) or MADV_FREE (which marks pages as freeable/garbage), MADV_COLD preserves the page contents [2][3]. If the pages are reclaimed, they will be swapped out or written back to backing storage, but they are not immediately discarded [3]. Applicability: The kernel may choose to ignore the hint for certain pages if it is not applicable [1]. For instance, implementation details often restrict the operation to pages that are not mapped by multiple processes (or have a low map count), because shared pages are harder to manage for targeted reclamation [5][3]. Relationship to Reclaim: MADV_COLD is effectively an "aging" hint [3]. By deactivating the pages, it ensures that when the system's background reclaim processes (such as kswapd) scan the LRU lists, these pages are encountered and selected for reclamation more quickly than they would have been otherwise [4]. Unlike MADV_PAGEOUT, which requests immediate reclamation, MADV_COLD allows the kernel to defer the actual work until memory pressure makes it necessary [2][6].
Citations:
- 1: https://man7.org/linux/man-pages/man2/madvise.2.html
- 2: https://lwn.net/Articles/790123/
- 3: https://d3s.mff.cuni.cz/files/teaching/nswi161/2023_24/08_linux_memory.pdf
- 4: https://lists.openwall.net/linux-kernel/2019/05/31/636
- 5: https://lists.openwall.net/linux-kernel/2019/06/03/210
- 6: https://lore.kernel.org/lkml/20190603053655.127730-1-minchan@kernel.org/
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files | rg '(^|/)s2_(mapped_file|pipeline)\.(cpp|h|hpp)$|(^|/)s2_mapped_file\.cpp$|(^|/)s2_pipeline\.cpp$' || true
echo
echo "== s2_mapped_file.cpp outline =="
ast-grep outline src/s2_mapped_file.cpp --view compact || true
echo
echo "== relevant mapped_file lines =="
cat -n src/s2_mapped_file.cpp | sed -n '1,280p'
echo
echo "== warm_page_cache usages =="
rg -n "warm_page_cache|drop_page_cache|MADV_COLD|MADV_WILLNEED|MADV_SEQUENTIAL" src package.xml CMakeLists.txt .github 2>/dev/null || trueRepository: rodrigomatta/s2.cpp
Length of output: 9017
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== pipeline relevant sections =="
cat -n src/s2_pipeline.cpp | sed -n '420,465p;t;q'
cat -n src/s2_pipeline.cpp | sed -n '1045,1075p;t;q'
cat -n src/s2_pipeline.cpp | sed -n '1100,1120p;t;q'
cat -n src/s2_pipeline.cpp | sed -n '1475,1495p;t&q'
echo
echo "== declarations =="
cat -n include/s2_pipeline.h | sed -n '1,180p'
cat -n include/s2_mapped_file.h | sed -n '1,140p'
echo
echo "== source text facts around warm_page_cache call =="
rg -n -C 8 "warm_page_cache|Slow-AR|deferr|defer" src/s2_pipeline.cpp include || trueRepository: rodrigomatta/s2.cpp
Length of output: 194
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== s2_pipeline.cpp relevant sections =="
cat -n src/s2_pipeline.cpp | sed -n '420,465p'
cat -n src/s2_pipeline.cpp | sed -n '1045,1075p'
cat -n src/s2_pipeline.cpp | sed -n '1100,1120p'
cat -n src/s2_pipeline.cpp | sed -n '1475,1495p'
echo
echo "== declarations =="
cat -n include/s2_pipeline.h | sed -n '1,180p'
cat -n include/s2_mapped_file.h | sed -n '1,140p'
echo
echo "== source text facts around warm_page_cache call =="
rg -n -C 8 "warm_page_cache|Slow-AR|deferr|defer" src/s2_pipeline.cpp include || trueRepository: rodrigomatta/s2.cpp
Length of output: 29769
Remove MADV_COLD from warm_page_cache().
warm_page_cache() is used to page in the deferred Slow-AR weights before the first request. MADV_COLD deactivates the same range and makes those pages easier reclaim targets, so the warmed pages can be dropped before use. Move cold reclamation to drop_page_cache() if needed.
♻️ Proposed fix
`#ifdef` __linux__
::madvise(data_, size_, MADV_WILLNEED);
::madvise(data_, size_, MADV_SEQUENTIAL);
-
-#ifdef MADV_COLD
- ::madvise(data_, size_, MADV_COLD);
-#endif
-
`#elif` defined(__APPLE__)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| #ifdef __linux__ | |
| ::madvise(data_, size_, MADV_WILLNEED); | |
| ::madvise(data_, size_, MADV_SEQUENTIAL); | |
| #ifdef MADV_COLD | |
| ::madvise(data_, size_, MADV_COLD); | |
| #endif | |
| `#ifdef` __linux__ | |
| ::madvise(data_, size_, MADV_WILLNEED); | |
| ::madvise(data_, size_, MADV_SEQUENTIAL); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_mapped_file.cpp` around lines 204 - 210, Remove the conditional
MADV_COLD madvise call from warm_page_cache(), leaving only the page-warming
advice there. If cold-page reclamation is required, apply MADV_COLD in
drop_page_cache() instead.
| const bool defer_weight_loading = params.enable_vram_swap && model().prefers_gpu(); | ||
|
|
||
| const auto codec_t1 = std::chrono::steady_clock::now(); | ||
| if (!defer_weight_loading) { | ||
| if (!model().allocate_and_load_weights()) { | ||
| safe_print_error_ln("Pipeline error: failed to allocate and load Slow-AR weights"); | ||
| return false; | ||
| } | ||
| } else { | ||
| safe_print_ln("[Pipeline] Deferring Slow-AR weight loading to first request (VRAM swap active)."); | ||
| model().mapped_file().warm_page_cache(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Deferral is decided from init params, but restoration is gated on request params.
Line 436 defers weight loading when params.enable_vram_swap is true at init time. Every restore path is then gated on the request copy of the same flag:
- offline: line 819
if (params.enable_vram_swap)wraps the phase-1 restore thread - streaming: line 1238
if (params.enable_vram_swap)wraps the streaming restore
PipelineParams is passed separately to init and to each synthesize_* call. If a caller initializes with VRAM swap enabled and then issues a request with enable_vram_swap = false, no restore runs and generate() executes against weight tensors whose data is still nullptr.
Gate the restore on the model state instead of the request flag. model().is_weights_on_gpu() and model().allocate_and_load_weights() already give you the needed check, and allocate_and_load_weights() returns early when the weights are present.
🛡️ Suggested direction
Add an unconditional guard at the top of synthesize_prompt_codes_locked and synthesize_streaming_prompt_codes_locked, before any generation:
if (!model().is_weights_on_gpu() && model_prefers_gpu_) {
model().acquire_compute_resources();
if (!model().restore_weights_to_gpu()) {
safe_print_error_ln("Pipeline error: Slow-AR weight restore failed.");
return false;
}
}Keep the existing enable_vram_swap blocks for the overlap and eviction optimizations only.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 436 - 446, Update
synthesize_prompt_codes_locked and synthesize_streaming_prompt_codes_locked to
restore weights based on model state, not the request’s enable_vram_swap flag:
before generation, when model().is_weights_on_gpu() is false and
model_prefers_gpu_ is true, acquire compute resources and call
restore_weights_to_gpu(), returning false with the existing error-reporting
style on failure. Retain the current enable_vram_swap blocks only for overlap
and eviction optimizations.
| const bool need_encoder_vram = params.enable_vram_swap && codec_prefers_gpu_; | ||
| if (need_encoder_vram) { | ||
| if (codec().is_decoder_on_gpu()) { | ||
| codec().free_decoder_weights(); | ||
| } | ||
| codec().restore_encoder_weights(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Weight-restore failures are discarded on every VRAM-swap path. AudioCodec::restore_encoder_weights, AudioCodec::restore_decoder_weights and SlowARModel::restore_weights_to_gpu all return bool, but these three sites drop the result and proceed straight into encode, decode, or generation. A failed restore leaves the tensors with no backing buffer, so the next ggml call dereferences null instead of reporting an error. GPU allocation failure is the exact condition this feature manages, so the path is reachable. Only the offline model restore at lines 824-828 checks its result today.
src/s2_pipeline.cpp#L546-L552: checkcodec().restore_encoder_weights()at line 551 and return false before thecodec().encode(...)call at line 555.src/s2_pipeline.cpp#L1083-L1089: checkcodec().restore_decoder_weights()at line 1086 and fail the request beforedecode_codes_windowedat line 1092.src/s2_pipeline.cpp#L1243-L1256: checkmodel().restore_weights_to_gpu()at line 1246 andcodec().restore_decoder_weights()at line 1255; report throughsink.on_error(...)and return false, matching the offline path at lines 824-828.
📍 Affects 1 file
src/s2_pipeline.cpp#L546-L552(this comment)src/s2_pipeline.cpp#L1083-L1089src/s2_pipeline.cpp#L1243-L1256
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 546 - 552, Weight-restore results are
ignored on three VRAM-swap paths. In src/s2_pipeline.cpp lines 546-552, check
codec().restore_encoder_weights() and return false before encode; in lines
1083-1089, check codec().restore_decoder_weights() and fail before
decode_codes_windowed; in lines 1243-1256, check both
model().restore_weights_to_gpu() and codec().restore_decoder_weights(), report
failures through sink.on_error(...), and return false, matching the offline
restore handling.
| std::thread vram_phase1_thread; | ||
| bool vram_phase1_ok = true; | ||
|
|
||
| if (params.enable_vram_swap) { | ||
| vram_phase1_thread = std::thread([this, ¶ms, &vram_phase1_ok]() { | ||
| if (model_prefers_gpu_ && !model().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Restoring Slow-AR to VRAM for generation..."); | ||
| model().acquire_compute_resources(); | ||
| if (!model().restore_weights_to_gpu()) { | ||
| safe_print_error_ln("Pipeline error: Slow-AR weight restore failed."); | ||
| vram_phase1_ok = false; | ||
| return; | ||
| } | ||
| safe_print_ln("[VRAM Diag] Post-SlowAR restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| if (!model_prefers_gpu_ && codec_prefers_gpu_ && !codec().is_decoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Pre-loading Audio Codec decoder to VRAM (hiding behind CPU gen)..."); | ||
| codec().restore_decoder_weights(); | ||
| safe_print_ln("[VRAM Diag] Post-Decoder restore: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| if (model_prefers_gpu_ && codec_prefers_gpu_) { | ||
| if (codec().is_encoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing codec encoder from VRAM (not needed during generation)..."); | ||
| codec().free_encoder_weights(); | ||
| } | ||
| if (codec().is_decoder_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing codec decoder from VRAM (not needed during generation)..."); | ||
| codec().free_decoder_weights(); | ||
| } | ||
| if (codec().is_weights_on_gpu()) { | ||
| safe_print_ln("[Pipeline] Freeing Audio Codec from VRAM for Slow-AR generation..."); | ||
| codec().free_gpu_weights(); | ||
| } | ||
| safe_print_ln("[VRAM Diag] Post-Codec free: Slow-AR=" + std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| } | ||
|
|
||
| safe_print_ln("[VRAM Diag] End-Phase1: Slow-AR=" + | ||
| std::to_string(model().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB, Codec=" + | ||
| std::to_string(codec().get_gpu_memory_usage_bytes() / 1024 / 1024) + " MB"); | ||
| }); | ||
| } | ||
|
|
||
| const int32_t num_codebooks = model().hparams().num_codebooks; | ||
| PromptTensor prompt = build_prompt( | ||
| tokenizer(), params.text, params.prompt_text, | ||
| ref_codes, | ||
| num_codebooks, T_prompt); | ||
|
|
||
| ref_codes, num_codebooks, T_prompt); | ||
| int32_t max_seq_len = prompt.cols + params.gen.max_new_tokens; | ||
|
|
||
| model().clear_kv_cache(); | ||
|
|
||
| const auto kv_t0 = std::chrono::steady_clock::now(); | ||
| if (!model().init_kv_cache(max_seq_len)) { | ||
| safe_print_error_ln("Pipeline error: init_kv_cache failed."); | ||
| std::thread kv_init_thread; | ||
| bool kv_init_ok = true; | ||
|
|
||
| kv_init_thread = std::thread([&]() { | ||
| kv_init_ok = model().init_kv_cache(max_seq_len); | ||
| }); | ||
|
|
||
| if (vram_phase1_thread.joinable()) { | ||
| vram_phase1_thread.join(); | ||
| } | ||
| if (!vram_phase1_ok) { | ||
| kv_init_thread.join(); | ||
| return false; | ||
| } | ||
| const auto kv_t1 = std::chrono::steady_clock::now(); | ||
|
|
||
| const auto gen_t0 = std::chrono::steady_clock::now(); | ||
| GenerateResult res = generate(model(), tokenizer().config(), prompt, params.gen); | ||
| const auto gen_t1 = std::chrono::steady_clock::now(); | ||
|
|
||
| if (res.n_frames == 0) { | ||
| safe_print_error_ln("Pipeline error: generation produced no frames."); | ||
| kv_init_thread.join(); | ||
| if (!kv_init_ok) { | ||
| safe_print_error_ln("Pipeline error: init_kv_cache failed."); | ||
| return false; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Data race: three threads mutate SlowARModel state at the same time.
The phase-1 thread starts at line 820 and the main thread keeps running. From line 866 a third thread joins in. All three touch the same SlowARModel and AudioCodec state, and the first join is only at line 876.
Concurrent accesses:
- Phase-1 thread, line 824:
restore_weights_to_gpu()→allocate_and_load_weights()→allocate_weight_buffers(), which rewritesweights_.model_bufs_gpu,weights_.model_bufs_cpu, and every tensor'sdataandbuffer, then callsacquire_compute_resources(), which writessched_andfast_sched_. - Main thread, line 866:
clear_kv_cache()freeskv_buf_andctx_kv_and nulls them. - KV thread, line 873:
init_kv_cache()writesctx_kv_,memory_k_,memory_v_,kv_buf_,max_seq_len_,n_past_, and allocates onbackend_gpu_.
Two concrete failures:
- Use-after-free. The diagnostics at lines 829, 835, 851 and 854-856 call
model().get_gpu_memory_usage_bytes(), which readskv_buf_(src/s2_model.cppline 1207) and iteratesweights_.model_bufs_gpu. Lines 866 and 873 free and reassignkv_buf_from other threads at the same time. - Concurrent backend allocation.
init_kv_cachecallsggml_backend_alloc_ctx_tensorsonbackend_gpu_while the phase-1 thread allocates GPU weight buffers and creates schedulers on the same backend. ggml backends do not support concurrent allocation on one device.
This is the same class of defect as the previously reported race around the hot-swap offload thread, but a different site and a different set of threads.
Serialize the two operations, or restrict the phase-1 thread to the weight restore and move every diagnostic read after both joins.
🔒 Suggested direction
- Remove every
get_gpu_memory_usage_bytes()call from the phase-1 lambda. Emit one diagnostic line after line 884, when both threads have joined. - Do not overlap
init_kv_cachewith the weight restore whenbackend_gpu_is shared. Either joinvram_phase1_threadbefore spawningkv_init_thread, or keep the overlap only when the KV cache is CPU-resident. - Move
model().clear_kv_cache()(line 866) before the phase-1 thread is spawned.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 816 - 888, Serialize VRAM phase-1 work with
KV-cache operations in the generation flow: move model().clear_kv_cache() before
starting vram_phase1_thread, and ensure vram_phase1_thread joins before spawning
kv_init_thread unless the KV cache is definitively CPU-resident. Remove all
get_gpu_memory_usage_bytes() calls from the phase-1 lambda, then emit the
combined VRAM diagnostic only after both threads have joined, while preserving
vram_phase1_ok and kv_init_ok failure handling.
| std::thread decode_thread([&]() { | ||
| int32_t last_committed = 0; | ||
| while (true) { | ||
| std::unique_lock<std::mutex> lock(decode_mtx); | ||
| decode_cv.wait(lock, [&]() { | ||
| return frames_available.load() > last_committed | ||
| || gen_done.load(); | ||
| }); | ||
| const int32_t avail = frames_available.load(); | ||
| const bool done = gen_done.load(); | ||
| lock.unlock(); | ||
| if (avail <= last_committed && done) | ||
| break; | ||
| if (!decode_window(avail, done)) { | ||
| decode_failed = true; | ||
| break; | ||
| } | ||
| last_committed = committed_frames; | ||
| } | ||
| }); | ||
|
|
||
| gen_params.on_frame = [&](const FrameCallbackData & fcd) -> bool { | ||
| { | ||
| std::lock_guard<std::mutex> lock(decode_mtx); | ||
| for (int32_t cb = 0; cb < fcd.num_codebooks; ++cb) | ||
| accum[cb].push_back(fcd.codes[cb]); | ||
| } | ||
| frames_available.store(fcd.total_frames); | ||
| decode_cv.notify_one(); | ||
| return true; | ||
| }; | ||
|
|
||
| const auto gen_t0 = std::chrono::steady_clock::now(); | ||
| res = generate(model(), tokenizer().config(), prompt, gen_params); | ||
| const auto gen_t1 = std::chrono::steady_clock::now(); | ||
| gen_ms = std::chrono::duration<double, std::milli>(gen_t1 - gen_t0).count(); | ||
|
|
||
| { | ||
| std::lock_guard<std::mutex> lock(decode_mtx); | ||
| gen_done.store(true); | ||
| } | ||
| decode_cv.notify_one(); | ||
| decode_thread.join(); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
decode_thread is not joined if generate throws.
decode_thread is created at line 979 and joined at line 1021. Between those lines, line 1012 calls generate(...), and lines 1000-1009 run the on_frame callback. If any of that throws, the std::thread destructor runs on a joinable thread and calls std::terminate.
The repository convention allows std::runtime_error for hard failures, so this is reachable. Wrap the thread in an RAII joiner.
🛡️ Proposed fix
Add a small joiner near CodecDecodeCacheScope (around line 119):
struct ThreadJoiner {
explicit ThreadJoiner(std::thread & t) : t_(t) {}
~ThreadJoiner() { if (t_.joinable()) t_.join(); }
std::thread & t_;
};Then guard the decode thread and drop the manual join:
std::thread decode_thread([&]() {
...
});
+ ThreadJoiner decode_thread_joiner(decode_thread); decode_cv.notify_one();
- decode_thread.join();
+ if (decode_thread.joinable()) decode_thread.join();
const auto decode_thread_t1 = std::chrono::steady_clock::now();Apply the same guard to vram_phase1_thread and kv_init_thread at lines 816-884.
As per coding guidelines: "Mix bool returns for recoverable operations and std::runtime_error for hard failures in C++".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 979 - 1021, Add an RAII thread-joining
helper near CodecDecodeCacheScope, then guard decode_thread, vram_phase1_thread,
and kv_init_thread with it so joinable threads are joined during exceptional
exits. Replace the manual decode_thread.join() with the helper’s cleanup while
preserving existing completion signaling and thread behavior.
Source: Coding guidelines
| if (params.enable_vram_swap && params.is_persistent && | ||
| params.enable_hot_swap && !params.more_segments_pending) { | ||
| safe_print_ln("[Pipeline] Hot-Swap: Releasing compute resources..."); | ||
| model().free_compute_buffers(); | ||
| safe_print_ln("[Pipeline] Hot-Swap: Spawning background thread to free VRAM & RAM..."); | ||
| std::thread offload_thread([this]() { | ||
| if (model().is_weights_on_gpu()) model().free_gpu_weights(); | ||
| if (codec().is_decoder_on_gpu()) codec().free_decoder_weights(); | ||
| if (codec().is_encoder_on_gpu()) codec().free_encoder_weights(); | ||
| if (codec().is_weights_on_gpu()) codec().free_gpu_weights(); | ||
| model().mapped_file().drop_page_cache(); | ||
| codec().mapped_file().drop_page_cache(); | ||
| safe_print_ln("[Pipeline] Hot-Swap: Background VRAM & RAM free complete."); | ||
| }); | ||
| pending_offload_thread_ = std::move(offload_thread); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
The offload thread still races the main-thread reads that follow it.
The thread spawned at line 1058 immediately calls model().free_gpu_weights() and codec().free_gpu_weights(), which clear weights_.model_bufs_gpu and reset residency flags. The main thread continues past line 1068 and reads the same state:
- line 1136:
model().clear_kv_cache() - lines 1170-1173:
model().get_gpu_memory_usage_bytes()andcodec().get_gpu_memory_usage_bytes(), which iteratemodel_bufs_gpuwhile the thread clears it
Move the Post-Phase3 diagnostic and clear_kv_cache() before the spawn, and make the spawn the last action in the function.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/s2_pipeline.cpp` around lines 1053 - 1068, In the hot-swap flow, move the
Post-Phase3 diagnostic and model().clear_kv_cache() before the background thread
is created, then make the offload-thread creation and assignment to
pending_offload_thread_ the final action in the function. Ensure no subsequent
code reads model or codec GPU residency, buffers, or memory usage after the
spawn.
1 participant
What this solves
In the final stages when it's time to process Audio Codec, Slow-AR isn't evicted from VRAM despite being no longer needed. So our VRAM footprint becomes
Slow-AR + compute buffers + Audio Codec + compute buffers + KV cachespeaking at6-7 GB on q8_0with transient spikes during buffer allocations (possibly due to fragmentaton?). In effect this causes high memory pressure and leads to OOMs if you were nearing hardware limit already.To solve this we selectively load the submodels and evict them from VRAM once they finish. By doing this we achieve 2.2 GB VRAM usage during Audio Codec phase of processing, and we also only load the Audio Codec when it's needed so the initial VRAM load is also lowered by ~1 GB.
With memory-mapped page cache sitting in RAM we can near-instantly fetch them when needed, and we hide the latency of this (mostly PCIe) fetch by scheduling it in background thread as we finish other sequential processes.
We've two atomic commits for easier review:
The new default behavior after the second commit can be categorized by opt-in and opt-out:
Opt out vram-swap
Unless you opt out, we evict already-processed submodels from VRAM, and lazily loads them as needed. Because we're using mmap these loads are near-instant since we fetch them from RAM instead of disk. We also keep compute buffers in VRAM as they're relatively small but take longer to re-initiate (~170 MB in vulkan). Your OS might (partially) reclaim the RAM page cache incurring some disk IO for the reclaimed pages, but it's generally as fast as keeping the models loaded in VRAM all the time. This depends on page faults. It's also more efficient than fread (double copy) since we don't load all tensors at once and use zero-copy DMA.
--no-vram-swapOpt in hot-swap
With the opt-in hot-swap behavior, we free anything from not only VRAM but also RAM. This includes compute buffers and the mmap page cache, landing us at a 100 MB RAM + 25 MB VRAM footprint. This is the "I don't trust my drivers & OS to free memory in a timely fashion and want to evict them aggressively, I need my VRAM back immediately for other tasks" option. This also ensures you'll need to read from disk each request, so I advise only using this if you absolutely need lowest memory footprint or the model is in m.2 SSD. It's a rather "dumb" solution compared to vram-swap. I advise against using this unless you know what you're doing as this can take 4x longer to process compared to vram-swap only. You're in effect doing cold boots every request and purely limited by your IO speed (mines a slow 200 MB/s SATA SSD)
--hot-swapKey Features
1. Zero-Copy
mmapArchitectureMappedFileRAII wrapper (POSIXmmap/ WindowsCreateFileMapping).2. Lazy Weight Allocation
encode()ordecode()is called viaensure_weights_loaded().mmappointer into system RAM.3. Phase-Gated VRAM Swapping
ggml_backend_synchronizeto nudge drivers to free VRAM instead of deferring & thrashing.4. Hot-Swap Mode (
--hot-swap)madvise(MADV_DONTNEED)/posix_fadvise(DONTNEED)to evict the 5.3 GB GGUF file from the OS page cache.Practical usage/footprint
--hot-swap--no-vram-swapSome metrics
Read me!:
--no-vram-swapflag.Notes
Summary by CodeRabbit
New Features
Performance